# --------------------------------------------------------
# SiamMask
# Licensed under The MIT License
# Written by Qiang Wang (wangqiang2015 at ia.ac.cn)
# --------------------------------------------------------
import glob
from tools.test import *
from experiments.siammask_sharp.custom import Custom

imgPathList = ['1_Car_01',
'1_Car_02',
'1_Car_03',
'1_Car_04',
'1_Car_05',
'1_Car_06',
'1_Car_07',
'1_Car_08',
'1_Car_09',
'1_Car_10',
'1_Car_11',
'1_Plane_01',
'1_Plane_02',
'1_Plane_03',
'1_Plane_04',
'1_Plane_05',
'1_Plane_06',
'1_Ship_01',
'1_Ship_02',
'1_Train_01']


siamModelPath = '../pth/SiamMask_DAVIS.pth'
imgPath = imgPathList[3]
track_mode = 'gt'
configPath = '../experiments/siammask_sharp/config.json'

parser = argparse.ArgumentParser(description='PyTorch Tracking Demo')
parser.add_argument('--config', dest='config', default=configPath,
                    help='hyper-parameter of SiamMask in json format')
args = parser.parse_args()

if __name__ == '__main__':
    # Setup device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    torch.backends.cudnn.benchmark = True
    # Setup Model
    cfg = load_config(args)

    siammask = Custom(anchors=cfg['anchors'])
    # if args.resume:
    #    assert isfile(args.resume), 'Please download {} first.'.format(args.resume)
    siammask = load_pretrain(siammask, siamModelPath)
    siammask.eval().to(device)

    # Parse Image file
    img_files = sorted(glob.glob(join('../img/' + imgPath + '/img', '*.jp*')))
    ims = [cv2.imread(imf) for imf in img_files]

    # Select ROI
    cv2.namedWindow("SiamMask", cv2.WND_PROP_FULLSCREEN)
    # cv2.setWindowProperty("SiamMask", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
    try:
        if track_mode == 'gt':
            gtpath = '../datasets/object/' + imgPath + '/groundtruth.txt'
            f = open(gtpath, 'r')
            x1, y1, w1, h1 = f.readline().split(',')
            x, y, w, h = int(x1), int(y1), int(w1), int(h1)
        else:
            init_rect = cv2.selectROI('SiamMask', ims[0], False, False)
            x, y, w, h = init_rect
            selpath = '../result/' + imgPath + '/groundtruthsel.txt'
            x, y, w, h = str(x), str(y), str(w), str(h)
            f = open(selpath, 'w')
            f.write(x+','+y+','+w+','+h)
            f.close()

        cv2.waitKey(100)
    except:
        exit()


    file1 = open('../result/' + imgPath + '/tgtPos.txt', 'w')

    toc = 0
    for f, im in enumerate(ims):
        tic = cv2.getTickCount()
        if f == 0:  # init
            x, y, w, h = int(x), int(y), int(w), int(h)
            target_pos = np.array([x + w / 2, y + h / 2])
            target_sz = np.array([w, h])
            state = siamese_init(
                im, target_pos, target_sz, siammask, cfg['hp'], device=device)  # init tracker
            x, y, w, h = str(x), str(y), str(w), str(h)
            file1.write(x + ',' + y + ',' + w + ',' + h + '\n')
        elif f > 0:  # tracking
            cv2.putText(im, str(f)+'FPS', (10, 20),
                        cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
            state = siamese_track(
                state, im, mask_enable=True, refine_enable=True, device=device)  # track
            location = state['ploygon'].flatten()
            mask = state['mask'] > state['p'].seg_thr

            im[:, :, 2] = (mask > 0) * 255 + (mask == 0) * im[:, :, 2]
            cv2.polylines(im, [np.int0(location).reshape(
                (-1, 1, 2))], True, (0, 255, 0), 3)
            cv2.imshow('SiamMask', im)
            key = cv2.waitKey(1)
            ######
            x, y = state['target_pos']
            w, h = state['target_sz']
            x = x - w/2
            y = y - h/2
            x, y, w, h = str(int(x)), str(int(y)), str(int(w)), str(int(h))
            if f != (len(ims)-1):
                file1.write(x + ',' + y + ',' + w + ',' + h + '\n')
            else:
                file1.write(x + ',' + y + ',' + w + ',' + h)

                ######
            if key > 0:
                break

        toc += cv2.getTickCount() - tic
    toc /= cv2.getTickFrequency()
    file1.close()
    fps = f / toc
    print('SiamMask Time: {:02.1f}s Speed: {:3.1f}fps (with visulization!)'.format(
        toc, fps))
